import math
import random

import numpy as np
import torch
from joblib import load
from sklearn.metrics.pairwise import cosine_similarity

from ETM import normalization

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sbert_result = load(r"data/20ng/data_20ng_SBERT_embedding.pkl")
etm_theta_result = load(r"data/20ng/data_20ng_etm_theta.pkl")
# print(sbert_result.keys())
sentence = sbert_result['doc_sentences']
sentences_vecs = sbert_result["doc_embeddings"]
sentence_label = sbert_result["all_label"]
doc_theta = etm_theta_result['theta']

train_len = sbert_result["train_len"]
sbt_vecs = sbert_result["doc_embeddings"][0:train_len]


def take_spl(spl_num):
    sbt_spls = []
    theta_spls = []
    _20ng_label = []
    spl_indxs = random.sample(range(0, len(sentences_vecs) + 1), spl_num)
    spl_indxs.sort()
    for i in range(0, spl_num):
        sbt_spls.append(sentences_vecs[spl_indxs[i]])
        theta_spls.append(doc_theta[spl_indxs[i]])
        _20ng_label.append(sentence_label[spl_indxs[i]])

    cos_sim_sbt = cosine_similarity(sbt_spls)
    cos_sim_theta = cosine_similarity(theta_spls)

    norm_mode = 'Z_ScoreNormalization'
    z_norm_sim_sbt = normalization.normalize_cos_sim(cos_sim_sbt, spl_indxs=spl_indxs, norm_mode=norm_mode)
    z_norm_sim_theta = normalization.normalize_cos_sim(cos_sim_theta, spl_indxs=spl_indxs, norm_mode=norm_mode)
    z_re = z_norm_sim_sbt - z_norm_sim_theta
    for i in range(z_re.size):
        z_re[i] = math.fabs(z_re[i])
    z_norm_sim_sbt_ = z_norm_sim_sbt[:, np.newaxis]
    z_norm_sim_theta_ = z_norm_sim_theta[:, np.newaxis]
    z_re_ = z_re[:, np.newaxis]

    norm_mode = 'MaxMinNormalization'
    m_norm_sim_sbt = normalization.normalize_cos_sim(cos_sim_sbt, spl_indxs=spl_indxs, norm_mode=norm_mode)
    m_norm_sim_theta = normalization.normalize_cos_sim(cos_sim_theta, spl_indxs=spl_indxs, norm_mode=norm_mode)
    m_re = m_norm_sim_sbt - m_norm_sim_theta
    for i in range(m_re.size):
        m_re[i] = math.fabs(m_re[i])
    m_norm_sim_sbt_ = m_norm_sim_sbt[:, np.newaxis]
    m_norm_sim_theta_ = m_norm_sim_theta[:, np.newaxis]
    m_re_ = m_re[:, np.newaxis]

    return z_norm_sim_sbt, z_norm_sim_theta, m_norm_sim_sbt, m_norm_sim_theta, _20ng_label, spl_indxs


take_spl(spl_num=10)
